import torch.nn as nn
from .element import *


class SubLayerConnection(nn.Module):
    def __init__(self, d_model, dropout_p=0.1):
        super().__init__()
        self.norm = LayerNorm(d_model)
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, data, sublayer):
        result = self.dropout(sublayer(self.norm(data))) + data
        return result


class EncoderLayer(nn.Module):
    def __init__(self, d_model, self_attn, feed_forward, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayers = clones(SubLayerConnection(d_model, dropout), 2)

    def forward(self, data, mask):
        data = self.sublayers[0](
            data, lambda x: self.self_attn(query=x, key=x, value=x, mask=mask)
        )
        data = self.sublayers[1](data, lambda x: self.feed_forward(x))
        return data


class Encoder(nn.Module):
    def __init__(self, layer: EncoderLayer, N):
        super().__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.d_model)

    def forward(self, data, mask=None):
        for layer in self.layers:
            data = layer(data, mask)
        return self.norm(data)
